/**
 * Copyright 2012 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.cluster.mom;

import java.util.Arrays;
import java.util.logging.Logger;

import com.google.common.base.Preconditions;

import edu.byu.nlp.data.SparseFeatureVector;
import edu.byu.nlp.math.Math2;
import edu.byu.nlp.util.DoubleArrays;
import edu.byu.nlp.util.Matrices;

/**
 * @author rah67
 *
 */
public class MoMParameters implements Cloneable {
	
	private static final Logger logger = Logger.getLogger(MoMParameters.class.getName());
	
	private double[] logPOfY;
	private double[][] logPOfXGivenY;
	
	// Holding space for computation of probabilities and related values.
	private double[] probs;
	
	private MoMParameters(double[] logPOfY, double[][] logPOfXGivenY) {
		this.logPOfY = logPOfY;
		this.logPOfXGivenY = logPOfXGivenY;
		this.probs = new double[logPOfY.length];
	}
	
	public double[] getLogPOfY() {
		return logPOfY;
	}
	
	public double[][] getLogPOfXGivenY() {
		return logPOfXGivenY;
	}
	
	/**
	 * Computes the logarithm of the joint probability of x and y, i.e.,
	 * log p(y, x) = log p(y) + \sum_v log p(v | y)^x_iv.
	 * Please note that this object maintains ownership of the returned array; the contents may change upon
	 * subsequent invocations of this method. Making a clone of the returned array is recommended. 
	 */
	public double[] logJoint(SparseFeatureVector instance) {
		// logJoint[y] = log p(y)
		System.arraycopy(logPOfY, 0, probs, 0, logPOfY.length);
		// The next call computes logPOfXGivenY x instance^T, i.e.,
		// logJoint[y] will become log p(y) + \sum_v log p(v | y)^x_iv
		instance.preMultiplyAsColumnAndAddTo(logPOfXGivenY, probs);
		return probs;
	}
	
	public double[] posterior(SparseFeatureVector instance) {
		double[] posterior = logJoint(instance);
		DoubleArrays.logNormalizeToSelf(posterior);
		DoubleArrays.expToSelf(posterior);
		return posterior;
	}
	
	public int getNumClasses() { 
		return logPOfY.length;
	}
	
	public int getNumFeatures() {
		return logPOfXGivenY[0].length;
	}
	
	public static MoMParameters fromLogProbabilities(double[] logPOfY, double[][] logPOfXGivenY,
			boolean copyParams, boolean checkDistribution) {
		Preconditions.checkArgument(logPOfY.length == logPOfXGivenY.length);
		
		if (copyParams) {
			logPOfY = logPOfY.clone();
			logPOfXGivenY = Matrices.clone(logPOfXGivenY);
		}
		
		if (checkDistribution) {
			logger.info("Checking distributions");
			if (!Math2.doubleEquals(DoubleArrays.logSum(logPOfY), 0.0, 1e-6))
				throw new IllegalArgumentException("logPOfY is not a proper distribution, logSum = " 
						+ DoubleArrays.logSum(logPOfY));
			for (int i = 0; i < logPOfXGivenY.length; i++) {
				if (!Math2.doubleEquals(DoubleArrays.logSum(logPOfXGivenY[i]), 0.0, 1e-6)) {
					throw new IllegalArgumentException("logPOfXGivenY is not a proper distribution, for y = " + i
							+ " logSum = " + DoubleArrays.logSum(logPOfXGivenY[i]));
				}
			}
		}
		return new MoMParameters(logPOfY, logPOfXGivenY);
	}
	
	public static MoMParameters fromProbabilities(double[] pOfY, double[][] pOfXGivenY,
			boolean copyParams, boolean checkDistribution) {
		Preconditions.checkArgument(pOfY.length == pOfXGivenY.length);
		
		if (copyParams) {
			pOfY = pOfY.clone();
			pOfXGivenY = Matrices.clone(pOfXGivenY);
		}

		DoubleArrays.logToSelf(pOfY);
		for (int i = 0; i < pOfXGivenY.length; i++) {
			DoubleArrays.logToSelf(pOfXGivenY[i]);
		}		
		
		return fromLogProbabilities(pOfY, pOfXGivenY, false, checkDistribution);
	}

	// TODO : add a check parameters flag which ensures the parameters are non-negative.
	public static MoMParameters fromWeights(double[] yCounts, double[][] xGivenYCounts, boolean copyParams) {
		
		if (copyParams) {
			yCounts = yCounts.clone();
			xGivenYCounts = Matrices.clone(xGivenYCounts);
		}
		
		DoubleArrays.normalizeToSelf(yCounts);
		for (int k = 0; k < xGivenYCounts.length; k++) {
			DoubleArrays.normalizeToSelf(xGivenYCounts[k]);
		}
		
		return fromProbabilities(yCounts, xGivenYCounts, false, false);
	}
	
	@Override
	public MoMParameters clone() {
		return MoMParameters.fromLogProbabilities(logPOfY, logPOfXGivenY, true, false);
	}

	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		result = prime * result + Arrays.deepHashCode(logPOfXGivenY);
		result = prime * result + Arrays.hashCode(logPOfY);
		return result;
	}

	@Override
	public boolean equals(Object obj) {
		if (this == obj)
			return true;
		if (obj == null)
			return false;
		if (getClass() != obj.getClass())
			return false;
		MoMParameters other = (MoMParameters) obj;
		if (!Arrays.equals(logPOfY, other.logPOfY))
			return false;
		if (!Arrays.deepEquals(logPOfXGivenY, other.logPOfXGivenY))
			return false;
		return true;
	}
	
}
